package api

import (
	"errors"
	"fmt"
	"wgmat/pkg/core"

	"github.com/cogentcore/webgpu/wgpu"
)

// Matrix 高阶矩阵API
type Matrix struct {
	ctx    *core.Context
	buffer core.Buffer
	rows   uint32
	cols   uint32
}

// NewMatrix 创建新矩阵（保持不变）
func NewMatrix(ctx *core.Context, data []float32, rows, cols uint32) (*Matrix, error) {
	expectedSize := int(rows * cols)
	if len(data) != expectedSize {
		return nil, fmt.Errorf("data size mismatch: expected %d elements, got %d", expectedSize, len(data))
	}

	shape := []uint32{rows, cols}
	usage := wgpu.BufferUsageStorage | wgpu.BufferUsageCopySrc | wgpu.BufferUsageCopyDst

	buffer, err := core.NewBuffer(ctx, data, shape, usage)
	if err != nil {
		return nil, fmt.Errorf("创建GPU buffer失败: %v", err)
	}

	return &Matrix{
		ctx:    ctx,
		buffer: buffer,
		rows:   rows,
		cols:   cols,
	}, nil
}

// NewMatrixFromBuffer 从现有buffer创建（保持不变）
func NewMatrixFromBuffer(ctx *core.Context, buffer core.Buffer, rows, cols uint32) *Matrix {
	return &Matrix{
		ctx:    ctx,
		buffer: buffer,
		rows:   rows,
		cols:   cols,
	}
}

// Add 创建矩阵加法命令（惰性求值）
// 返回Command对象和结果矩阵，需要后续Execute执行
// 示例:
//
//	cmd, result := matA.Add(matB)
//	executor.Execute(cmd) // 同步执行
func (m *Matrix) Add(other *Matrix) (core.Command, *Matrix, error) {
	if m.rows != other.rows || m.cols != other.cols {
		return nil, nil, fmt.Errorf("矩阵维度不匹配: A=%dx%d, B=%dx%d",
			m.rows, m.cols, other.rows, other.cols)
	}

	// 验证buffer usage
	if m.buffer.Raw().GetUsage()&wgpu.BufferUsageStorage == 0 {
		return nil, nil, fmt.Errorf("矩阵A buffer缺少Storage usage (当前: 0x%x, 需要: 0x%x)",
			m.buffer.Raw().GetUsage(), wgpu.BufferUsageStorage)
	}
	if other.buffer.Raw().GetUsage()&wgpu.BufferUsageStorage == 0 {
		return nil, nil, fmt.Errorf("矩阵B buffer缺少Storage usage (当前: 0x%x, 需要: 0x%x)",
			other.buffer.Raw().GetUsage(), wgpu.BufferUsageStorage)
	}

	// 创建结果buffer
	usage := wgpu.BufferUsageStorage | wgpu.BufferUsageCopySrc
	resultBuffer, err := core.NewBuffer(m.ctx, nil, []uint32{m.rows, m.cols}, usage)
	if err != nil {
		return nil, nil, fmt.Errorf("创建结果buffer失败: %v", err)
	}

	// 创建结果矩阵
	result := NewMatrixFromBuffer(m.ctx, resultBuffer, m.rows, m.cols)

	// 创建命令对象（不立即执行）
	cmd := &MatrixAddCommand{
		ctx:  m.ctx,
		a:    m.buffer,
		b:    other.buffer,
		c:    resultBuffer,
		rows: m.rows,
		cols: m.cols,
	}

	return cmd, result, nil
}

// Mul 创建矩阵乘法命令（演示可扩展性）
// 返回Command对象和结果矩阵
func (m *Matrix) Mul(other *Matrix) (core.Command, *Matrix, error) {
	if m.cols != other.rows {
		return nil, nil, fmt.Errorf("矩阵维度不匹配: A=%dx%d, B=%dx%d (A.cols != B.rows)", m.rows, m.cols, other.rows, other.cols)
	}

	usage := wgpu.BufferUsageStorage | wgpu.BufferUsageCopySrc
	resultBuffer, err := core.NewBuffer(m.ctx, nil, []uint32{m.rows, other.cols}, usage)
	if err != nil {
		return nil, nil, fmt.Errorf("创建结果buffer失败: %v", err)
	}
	result := NewMatrixFromBuffer(m.ctx, resultBuffer, m.rows, other.cols)

	// 根据规模选择策略
	totalElements := m.rows * m.cols * other.cols
	if totalElements > 128*128*128 {
		// ✅ 中等规模：16x16（标准tile大小）
		cmd := &MatrixTiledMulCommand{
			ctx:   m.ctx,
			a:     m.buffer,
			b:     other.buffer,
			c:     resultBuffer,
			rowsA: m.rows,
			colsA: m.cols,
			colsB: other.cols,
			tileX: 16,
			tileY: 16,
		}
		return cmd, result, nil
	} else {
		// 小规模：保留原始 naive kernel
		cmd := &MatrixMulCommand{
			ctx:   m.ctx,
			a:     m.buffer,
			b:     other.buffer,
			c:     resultBuffer,
			rowsA: m.rows,
			colsA: m.cols,
			colsB: other.cols,
		}
		return cmd, result, nil
	}
}

// QR 执行矩阵的QR分解
// 返回 Q (m x k), R (m x n)，其中 k = min(m, n) for economy mode
// 或者 k = m for full mode
// 示例:
//
//	cmd, Q, R, err := mat.QR(false) // economy mode
//	executor.Execute(cmd)
func (m *Matrix) QR(fullMode bool) (core.Command, *Matrix, *Matrix, error) {
	mRows, nCols := m.Shape()
	if mRows == 0 || nCols == 0 {
		return nil, nil, nil, fmt.Errorf("矩阵维度无效")
	}

	// 计算 Q 的列数 k
	var k uint32
	if fullMode {
		k = mRows // full mode: Q is m × m
	} else {
		k = nCols // economy mode: Q is m × n (assume m >= n; if not, still valid)
		if mRows < nCols {
			k = mRows // economy with m < n → Q is m × m
		}
	}

	device := m.ctx.Device()
	queue := m.ctx.Queue()

	// 1. 创建 R = A（深拷贝）
	rUsage := wgpu.BufferUsageStorage | wgpu.BufferUsageCopySrc
	rBuffer, err := core.NewBuffer(m.ctx, nil, []uint32{mRows, nCols}, rUsage)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("创建 R 矩阵失败: %v", err)
	}

	// 同步拷贝 A → R
	{
		enc, err := device.CreateCommandEncoder(nil)
		if err != nil {
			rBuffer.Release()
			return nil, nil, nil, fmt.Errorf("创建 encoder 失败: %v", err)
		}
		enc.CopyBufferToBuffer(m.buffer.Raw(), 0, rBuffer.Raw(), 0, m.buffer.Size())
		cmdBuf, err := enc.Finish(nil)
		enc.Release()
		if err != nil {
			rBuffer.Release()
			return nil, nil, nil, fmt.Errorf("finish encoder 失败: %v", err)
		}
		queue.Submit(cmdBuf)
		cmdBuf.Release()

		// 等待拷贝完成（确保 R 初始化完成）
		done := make(chan struct{})
		queue.OnSubmittedWorkDone(func(status wgpu.QueueWorkDoneStatus) {
			close(done)
		})
		<-done
	}

	// 2. 创建 Q 并初始化为单位矩阵（m × k）
	qUsage := wgpu.BufferUsageStorage | wgpu.BufferUsageCopySrc
	qData := make([]float32, mRows*k)
	for i := uint32(0); i < mRows && i < k; i++ {
		qData[i*k+i] = 1.0
	}

	qBuffer, err := core.NewBuffer(m.ctx, qData, []uint32{mRows, k}, qUsage)
	if err != nil {
		rBuffer.Release()
		return nil, nil, nil, fmt.Errorf("初始化 Q 矩阵失败: %v", err)
	}

	// 创建命令对象
	cmd := &MatrixQRCommand{
		ctx: m.ctx,
		q:   qBuffer,
		r:   rBuffer,
		m:   mRows,
		n:   nCols,
		k:   k,
	}

	return cmd,
		NewMatrixFromBuffer(m.ctx, qBuffer, mRows, k),
		NewMatrixFromBuffer(m.ctx, rBuffer, mRows, nCols),
		nil
}

// FetchAsync 异步读取矩阵数据（保持不变）
func (m *Matrix) FetchAsync(callback func([]float32, error)) {
	if asyncBuffer, ok := m.buffer.(interface {
		ReadAsync(callback func([]float32, error))
	}); ok {
		asyncBuffer.ReadAsync(callback)
		return
	}
	callback(nil, errors.New("buffer does not support async read"))
}

// Shape 返回矩阵维度
func (m *Matrix) Shape() (rows, cols uint32) {
	return m.rows, m.cols
}

// Release 释放资源
func (m *Matrix) Release() {
	if m.buffer != nil {
		m.buffer.Release()
	}
}
